import torch
from torchvision import transforms
from model import GoogLeNet, Inception
from PIL import Image


if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = GoogLeNet(Inception)
    model = model.to(device)
    # 加载模型参数
    model.load_state_dict(torch.load('checkpoints/classification_epoch_5.pth', map_location='cpu'))
    image = Image.open('images.jpg')
    normalize = transforms.Normalize([0.486, 0.453, 0.415], [0.069, 0.065, 0.067])
    transform = transforms.Compose([transforms.Resize((224, 224)),
                                    transforms.ToTensor(),
                                    normalize])
    image = transform(image)

    # 添加批次维度
    image = image.unsqueeze(0)
    classes = ['猫', '狗']
    with torch.no_grad():
        model.eval()
        image = image.to(device)
        output = model(image)
        pre_lab = torch.argmax(output, dim=1)
        result = pre_lab.item()
    print(f"预测值：{classes[result]}")




